# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: whpark

from coin.tool.slurm.slurm import SlurmArrayJob, SlurmJob


class JobGroup(object):
  def __init__(self, name, jobs, array_by_memories=None, max_active_tasks=None):
    self.name = name
    self.start_trigger = SlurmJob("%s_start_trigger" % name, 'echo 0')
    self.end_trigger = SlurmJob("%s_end_trigger" % name, 'echo 0')
    self.jobs = jobs
    self.is_skip = False

    if not self.jobs:
      self.end_trigger.set_afterok(self.start_trigger)
    else:
      if array_by_memories is not None:
        self.make_array(array_by_memories, max_active_tasks)

      self.end_trigger.set_afterok(self.jobs)
      for job in self.jobs:
        job.set_afterok(self.start_trigger)

  def get_start_trigger_job(self):
    return self.start_trigger

  def get_end_trigger_job(self):
    return self.end_trigger

  def set_afterok(self, job):
    if isinstance(job, JobGroup):
      return self.start_trigger.set_afterok(job.get_end_trigger_job())
    if isinstance(job, list):
      for a_job in job:
        assert isinstance(a_job, SlurmJob)
    else:
      assert isinstance(job, SlurmJob)
    return self.start_trigger.set_afterok(job)

  def get_all_jobs(self):
    if self.is_skip:
      return [self.start_trigger, self.end_trigger]
    return [self.start_trigger, self.end_trigger] + self.jobs

  def skip(self):
    self.end_trigger.set_afterok(self.start_trigger)
    self.is_skip = True

  def make_array(self, memories, max_active_tasks):
    assert len(set([job.cpu for job in self.jobs])) == 1
    cpu = self.jobs[0].cpu

    array_jobs = []
    prev_mem = 0
    for mem in sorted(memories):
      jobs = [job for job in self.jobs if job.mem > prev_mem and job.mem <= mem]
      if jobs:
        job = SlurmArrayJob("%s_%s_arr" % (self.name, mem),
                            jobs,
                            mem=mem,
                            cpu=cpu,
                            max_active_tasks=max_active_tasks)
        array_jobs.append(job)
      prev_mem = mem

    self.jobs = [job for job in self.jobs if job.mem > prev_mem] +\
        array_jobs
